import binascii
from gmssl.sm4 import CryptSM4, SM4_ENCRYPT, SM4_DECRYPT
import json


def hex_str_to_bytes(key_hex_str):
    """
    将十六进制格式的字符串转换为16字节的字节类型数据
    :param key_hex_str: 十六进制格式的字符串
    :return: 字节类型数据
    """
    return binascii.unhexlify(key_hex_str)


def check_key_iv_len(key, iv):
    if isinstance(key, str):
        key = hex_str_to_bytes(key)
    assert len(key) == 16, "密钥长度必须为16字节"
    assert isinstance(iv, bytes) and len(iv) == 16, "初始向量长度必须为16字节"


def sm4_encrypt_dict(key, input_dict, iv):
    """
    使用SM4算法（CBC模式）加密字典数据
    :param key: 加密密钥，支持十六进制格式字符串或字节类型，长度对应16字节
    :param input_dict: 需要加密的字典数据
    :param iv: 初始向量，字节类型，长度必须为16字节
    :return: 加密后的字节数据
    """
    if isinstance(key, str):
        key = hex_str_to_bytes(key)
    check_key_iv_len(key, iv)
    # 将字典序列化为字符串后再编码为字节数据
    data_to_encrypt = json.dumps(input_dict).encode('utf-8')
    crypt_sm4 = CryptSM4(mode=SM4_ENCRYPT)
    crypt_sm4.set_key(key, SM4_ENCRYPT)
    return crypt_sm4.crypt_cbc(iv, data_to_encrypt)


def sm4_decrypt_dict(key, encrypted_data, iv):
    """
    使用SM4算法（CBC模式）解密字节数据并还原为字典
    :param key: 解密密钥，支持十六进制格式字符串或字节类型，长度对应16字节
    :param encrypted_data: 加密后的字节数据
    :param iv: 初始向量，字节类型，长度必须为16字节
    :return: 解密并还原后的字典数据，如果还原失败返回None
    """
    if isinstance(key, str):
        key = hex_str_to_bytes(key)
    check_key_iv_len(key, iv)
    crypt_sm4 = CryptSM4(mode=SM4_DECRYPT)
    crypt_sm4.set_key(key, SM4_DECRYPT)
    decrypted_data = crypt_sm4.crypt_cbc(iv, encrypted_data)
    try:
        # 尝试将解密后的字节数据解码并反序列化为字典
        return json.loads(decrypted_data.decode('utf-8'))
    except json.JSONDecodeError:
        print("解密后的数据无法正确还原为字典")
        return None


def sm4_encrypt_text(key, text, iv):
    """
    使用SM4算法（CBC模式）加密文本数据
    :param key: 加密密钥，支持十六进制格式字符串或字节类型，长度对应16字节
    :param text: 需要加密的文本字符串
    :param iv: 初始向量，字节类型，长度必须为16字节
    :return: 加密后的字节数据
    """
    if isinstance(key, str):
        key = hex_str_to_bytes(key)
    check_key_iv_len(key, iv)
    crypt_sm4 = CryptSM4(mode=SM4_ENCRYPT)
    crypt_sm4.set_key(key, SM4_ENCRYPT)
    return crypt_sm4.crypt_cbc(iv, text.encode('utf-8'))


def sm4_decrypt_text(key, encrypted_text, iv):
    """
    使用SM4算法（CBC模式）解密字节数据并还原为文本
    :param key: 解密密钥，支持十六进制格式字符串或字节类型，长度对应16字节
    :param encrypted_text: 加密后的字节数据
    :param iv: 初始向量，字节类型，长度必须为16字节
    :return: 解密并还原后的文本字符串，如果还原失败返回None
    """
    if isinstance(key, str):
        key = hex_str_to_bytes(key)
    check_key_iv_len(key, iv)
    crypt_sm4 = CryptSM4(mode=SM4_DECRYPT)
    crypt_sm4.set_key(key, SM4_DECRYPT)
    decrypted_data = crypt_sm4.crypt_cbc(iv, encrypted_text)
    try:
        return decrypted_data.decode('utf-8')
    except UnicodeDecodeError:
        print("解密后的数据无法正确解码为文本")
        return None


def sm2_private_key_encrypt(key, private_key, iv):
    """
    使用SM4算法（CBC模式）加密SM2私钥（这里只是示例，实际情况私钥保密要求极高）
    :param key: SM4加密使用的16字节密钥，支持十六进制格式字符串或字节类型，用于加密SM2私钥
    :param private_key: SM2私钥，十六进制字符串形式
    :param iv: 初始向量，字节类型，长度必须为16字节
    :return: 加密后的结果，字节类型
    """
    if isinstance(key, str):
        key = hex_str_to_bytes(key)
    check_key_iv_len(key, iv)
    # 将十六进制字符串形式的私钥转换为字节类型（按十六进制解析）
    private_key_bytes = binascii.unhexlify(private_key)
    crypt_sm4 = CryptSM4(mode=SM4_ENCRYPT)
    crypt_sm4.set_key(key, SM4_ENCRYPT)
    return crypt_sm4.crypt_cbc(iv, private_key_bytes)


def sm2_private_key_decrypt(key, encrypted_private_key, iv):
    """
    使用SM4算法（CBC模式）解密加密后的SM2私钥并还原
    :param key: SM4解密使用的16字节密钥，支持十六进制格式字符串或字节类型，用于解密SM2私钥
    :param encrypted_private_key: 已经加密的SM2私钥，字节类型
    :param iv: 初始向量，字节类型，长度必须为16字节
    :return: 解密后的SM2私钥，字节类型，如果解密失败返回None
    """
    if isinstance(key, str):
        key = hex_str_to_bytes(key)
    check_key_iv_len(key, iv)
    crypt_sm4 = CryptSM4(mode=SM4_DECRYPT)
    crypt_sm4.set_key(key, SM4_DECRYPT)
    try:
        decrypted_bytes = crypt_sm4.crypt_cbc(iv, encrypted_private_key)
        # 将解密后的字节数据转换为十六进制字符串形式（方便查看和后续可能的对比等操作）
        return binascii.hexlify(decrypted_bytes).decode('utf-8')
    except:
        print("解密SM2私钥出现异常")
        return None


if __name__ == "__main__":
    # 两种形式的key均可使用
    # encrypt_key = '1234567890abcdef1234567890abcdef'
    encrypt_key = b'1234567890abcdef'
    # 定义16字节的初始向量（IV），实际应用中IV要保证随机性且妥善保管
    iv = b'abcdef0123456789'
    # 测试文本加密解密
    test_text = "这是一段需要加密的测试文本"
    encrypted_text_result = sm4_encrypt_text(encrypt_key, test_text, iv)
    print("加密后的文本结果:", encrypted_text_result)
    decrypted_text_result = sm4_decrypt_text(encrypt_key, encrypted_text_result, iv)
    if decrypted_text_result:
        print("解密并还原后的文本:", decrypted_text_result)

    # 测试字典加密解密
    test_dict = {
        "name": "张三",
        "age": 25,
        "hobbies": ["reading", "swimming"]
    }
    encrypted_dict_result = sm4_encrypt_dict(encrypt_key, test_dict, iv)
    print("加密后的字典结果:", encrypted_dict_result)
    decrypted_dict_result = sm4_decrypt_dict(encrypt_key, encrypted_dict_result, iv)
    if decrypted_dict_result:
        print("解密并还原后的字典:", decrypted_dict_result)

    # 模拟生成一个SM2私钥（实际应用中私钥来源有严格规范和安全要求）
    private_key = '00B9AB0B828FF68872F21A837FC303668428DEA11DCD1B24429D0C99E24EED83D5'
    encrypted_private_key_result = sm2_private_key_encrypt(encrypt_key, private_key, iv)
    print("加密后的SM2私钥:", encrypted_private_key_result)
    decrypted_private_key_result = sm2_private_key_decrypt(encrypt_key, encrypted_private_key_result, iv)
    if decrypted_private_key_result:
        print("解密后的SM2私钥:", decrypted_private_key_result)